#!/usr/bin/env python
#
# testing for keep-alive connection
#   agent-pool support
#

import time,socket,base64,fcntl,errno,sys
import threading
import Queue as queue

import celtuce.net.ioloop as ioloop
import celtuce.net.iostream as iostream
import celtuce.help.codec as codec

from celtuce.event import Event,ROUTE_GET,ROUTE_DIGEST,ROUTE_DEMAND,OP_RESPONSE,OP_FAILURE
from celtuce.event import OP_INSERT,OP_ERASE,OP_FIND

from celtuce.gatepacket import Packet
from celtuce.helper import getTimestamp
from celtuce.help import logging

from bisect import bisect_left,insort


URL = 'tcp://127.0.0.1:14300'

#######################################################################


class   ServerThread(threading.Thread) :
    def __init__(self, meta_url, tcp_port=0) :
        threading.Thread.__init__(self) 

        self._meta_url = meta_url
        self._tcp_server = None

        self._io_loop = ioloop.IOLoop()
        
        self._agent_pool = AgentPool(self._meta_url, self._io_loop)

        if tcp_port :
            self._tcp_port = tcp_port
            self._tcp_server = TcpServer(self._agent_pool.handleClientRequest, self._io_loop)

        self._logger = logging.getLogger('serverthread')
        ## __init__()

    def run(self) :
        self._tcp_server.listen( self._tcp_port )

        self._io_loop.start()
        ## run()

    def stop(self) :
        self._logger.info('stop IOLoop ...')
        self._io_loop.stop()
        ## stop()

    ## class ServerThread


class   TcpServer(object) :
    def __init__(self, request_callback, io_loop) :
        self._request_callback = request_callback
        self._io_loop = io_loop

        self._socket = None
        self._connections = {}
        self._logger = logging.getLogger('tcpserver')
        ## __init__()

    def listen(self, port, address='') :
        assert not self._socket
        self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)
        flags = fcntl.fcntl(self._socket.fileno(), fcntl.F_GETFD)
        flags |= fcntl.FD_CLOEXEC
        fcntl.fcntl(self._socket.fileno(), fcntl.F_SETFD, flags)
        self._socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self._socket.setblocking(0)
        self._socket.bind((address, port))
        self._socket.listen(128)
        self._io_loop.add_handler(self._socket.fileno(), self._handleEvents, self._io_loop.READ)
        ## listen()

    def _handleEvents(self, fd, events):
        while True:
            try:
                connection, address = self._socket.accept()
            except socket.error, e:
                if e[0] in (errno.EWOULDBLOCK, errno.EAGAIN):
                    return
                raise
            try:
                stream = iostream.IOStream(connection, io_loop=self._io_loop)
                self._connections[ address ] = TcpConnection(stream, address, self._request_callback)
            except:
                self._logger.excpt()
        ## _handleEvents()


    ## class TcpServer

class   TcpConnection(object) :
    def __init__(self, stream, address, request_callback) :
        self._stream = stream
        self._address = address
        self._request_callback = request_callback

        self._END_TAG = "\r\n\r\n"
        self._request_finished = False
        self._stream.read_until(self._END_TAG, self.__onRequest)

        self._keep_alive = True

        self._logger = logging.getLogger('tcpconnect')
        self._logger.debug('create TcpConnection for %s', self._address)
        ## __init__()

    def __del__(self) :
        #if self._stream.closed() :
        #    self._logger.debug('stream is closed when destroying')
        #else :
        #    self._logger.debug('stream is opened when destroying')
        self._logger.debug('destroy TcpConnection for %s', self._address)
        pass
        ## __del__()

    def __onRequest(self, data) :
        request = data
        #self._logger.debug('recv request  : %s', repr(request))
        #response = self._request_callback(request)
        self._request_callback(self, request)
        ## __onRequest()

    def sendResponse(self, response) :
        self.write(response)
        ## sendResponse()

    def write(self, data) :
        if data :
            #self._logger.debug('send response : %s', repr(data))
            self._stream.write(''.join((data, self._END_TAG)), self.__onWriteComplete)
        else :
            if not self._keep_alive :
                self.finish()
        ## write()

    def finish(self):
        if self._stream and not self._stream.writing():
            self.__finishRequest()

    def __onWriteComplete(self):
        if self._keep_alive :
            self._request_finished = False
            self._stream.read_until(self._END_TAG, self.__onRequest)
        else :
            self.__finishRequest()
        ## __onWriteComplete()

    def __finishRequest(self) :
        #self._logger.debug('finish request ...')
        self._stream.close()
        self._stream = None
        ## __finishRequest()

    ## class TcpConnection


#######################################################################

try :
    from hashlib import md5 as MD5
except :
    from md5 import md5 as MD5
import celtuce.help.url as urlparser


class   AgentPool(object) :

    def __init__(self, meta_url, io_loop) :
        self._io_loop = io_loop

        # such as tcp://host:port
        self._meta_url = meta_url
        parser = urlparser.URL(meta_url)
        self._meta_protocol = parser.protocol
        self._meta_host = parser.host
        self._meta_port = parser.port

        self._END_TAG = "\r\n\r\n"

        self._agents = {}
        self._requests = {}

        self._route_tablet = {}
        self._route_tablet_text = None
        self._route_tablet_digest = None
        self._route_last_update = 0

        self._id_base = id(self)
        self._id_serial = 1 
        self._id_maximum = 1024 * 1024 * 1024

        deadline = time.time() + 0.5
        self._io_loop.add_timeout( deadline, self.onUpdateRouteTable )

        self._when_to_update_route_table = 1.0

        self._logger = logging.getLogger('agentpool')
        ## __init__()

    def handleClientRequest(self, client_connection, request) :
        # STEP. parse packet, to get OP,Key,timestamp,etc ...
        packet = Packet.depack(request)
        if not packet :
            self._logger.error('fail to decode packet %r, just discard it.', request)
            return

        # STEP. lookup agent for this request
        agents = self.lookupAgent( packet )
        if not agents :
            self._logger.error('fail to lookup agent for (table:%r key:%r)', packet.table, packet.key)
            response = Packet()
            response.setResponse()
            response.respid = packet.packid
            response.table,response.op,response.key = packet.table,packet.op,packet.key
            response.state = Packet.RC_BAD_ROUTE
            response.error = 'fail to lookup agent for (table:%r key:%r)' % (packet.table, packet.key)
            self.sendClientResponse(client_connection, response)
            return

        # STEP. put request into request-pool
        timestamp = getTimestamp()
        result = {}
        for agent in agents :
            result[ agent ] = None
        self._requests[ (client_connection, packet, timestamp) ] = result

        # STEP. transform request into event
        evt = self.__transformRequest(packet, timestamp)

        # STEP. put request into queue of the agent
        for agent in agents :
            agent.sendRequest( client_connection, packet, timestamp, evt )

        ## handleClientRequest()

    def sendClientResponse(self, client_connection, response) :
        response_str,response_id = Packet.enpack(response)
        client_connection.sendResponse( response_str )
        ## sendClientResponse()

    def __transformRequest(self, packet, timestamp) :
        evt = Event()
        evt.request_id = self.genRequestID()

        if Packet.OP_INSERT == packet.op :
            evt.op = OP_INSERT
            evt.body = (packet.table, packet.key, packet.data, timestamp)
        elif Packet.OP_ERASE == packet.op :
            evt.op = OP_ERASE
            evt.body = (packet.table, packet.key, timestamp)
        elif Packet.OP_FIND == packet.op :
            evt.op = OP_FIND
            evt.body = (packet.table, packet.key, timestamp)

        else :
            self._logger.error('WHY not transform packet %s ??', packet)

        return evt
        ## __transformRequest()

    def feedResponse(self, client_connection, request, timestamp, agent, response) :
        id = (client_connection, request, timestamp)
        if id not in self._requests :
            # just discard some out-of-date response
            self._logger.error('response for (%s, %s, %s) is out of date', str(client_connection), str(request), str(timestamp))
            return

        result_set = self._requests[ id ]
        if agent not in result_set :
            self._logger.error('unknown agent for (%s, %s, %s)', str(client_connection), str(request), str(timestamp))
            return

        result_set[ agent ] = response

        # check that all response is returned
        want_more_response = False
        for _agent,_response in result_set.iteritems() :
            if _response is None :
                want_more_response = True
                break
        if want_more_response :
            return

        # merge all responses, and return
        done_set = []
        fail_set = []
        for _agent,_response in result_set.iteritems() :
            r = _response.body
            if r.isOk() :
                done_set.append( (_agent, r) )
            else :
                fail_set.append( (_agent, r) )

        # TODO : to find the optimal result
        if done_set :
            # success
            response_packet = self.__makeSuccessResponsePacket(done_set, request)
        else :
            response_packet = self.__makeFailureResponsePacket(request)
            # failure

        # handle error response
        if fail_set :
            self.__handleErrorResponses(done_set, fail_set)

        # free request slot
        del self._requests[ id ]

        response_str,packet_id = Packet.enpack( response_packet )
        client_connection.sendResponse( response_str )
        ## feedResponse()

    def __makeSuccessResponsePacket(self, done_set, request) :
        packet = Packet()
        packet.setResponse()
        packet.respid = request.packid
        packet.state = Packet.RC_OK
        packet.table,packet.op,packet.key = request.table, request.op, request.key

        if packet.op not in (Packet.OP_FIND,) :
            # for insert/erase operations, just return
            return packet

        # merge the results
        if len(done_set) == 1 :
            for _agent, _response in done_set :
                if isinstance(_response.body, dict) :
                    _key,_ts,_op,_data = [ _response.body[n] for n in range(len(_response.body)) ]
                else :
                    _key,_ts,_op,_data = _response.body
                packet.data = _data
                break
        else :
            result_set = []
            for _agent, _response in done_set :
                if isinstance(_response.body, dict) :
                    _key,_ts,_op,_data = [ _response.body[n] for n in range(len(_response.body)) ]
                else :
                    _key,_ts,_op,_data = _response.body
                insort(result_set, (_ts, _op, _data))
            _record = result_set[ -1 ]
            if 0 == _record[1] :
                # the last operation is insert
                packet.data = _record[2]

        return packet
        ## __makeSuccessResponsePacket()

    def __makeFailureResponsePacket(self, request) :
        packet = Packet()
        packet.setResponse()
        packet.respid = request.packid
        packet.state = Packet.RC_ERROR
        packet.table,packet.op,packet.key = request.table, request.op, request.key

        return packet
        ## __makeFailureResponsePacket()

    def __handleErrorResponses(self, done_set, fail_set) :
        for _agent,r in fail_set :
            self._logger.error('Error Response : %s', r.error)
        ## __handleErrorResponses()

    ####-----------------------------------------------------------####

    def onUpdateRouteTable(self) :
        self._logger.debug('update route table ...')

        try :
            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)
            s.connect( (self._meta_host, self._meta_port) )
        except :
            self._logger.excpt()
            self.onEndUpdateRouteTable()
            return

        request = Event()
        request.request_id = self.genRequestID()
        request.op = ROUTE_GET

        stream = iostream.IOStream(s, io_loop=self._io_loop)
        MetaConnection(stream, (self._meta_host, self._meta_port), request, self.updateRouteTable, self.onEndUpdateRouteTable)
        ## OnUpdateRouteTable()

    def onEndUpdateRouteTable(self) :
        deadline = time.time() + self._when_to_update_route_table
        self._io_loop.add_timeout( deadline, self.onUpdateRouteTable )
        ## OnEndUpdateRouteTable()

    def updateRouteTable(self, response) :
        if not response : return
        if isinstance(response.body, dict) :
            response.body = [ response.body[n] for n in range(len(response.body)) ]
        self._logger.debug('update route-table with %s', response.body)
        data,digest = response.body
        if not data or not digest : return
        
        if self._route_tablet_digest == digest :
            # got the same digest
            if self._route_tablet :
                self._when_to_update_route_table = 10.0
            else :
                self._when_to_update_route_table = 1.0
        else :
            # got new route data, update it
            self._route_tablet_text = data
            self._route_tablet_digest = digest
            self._route_last_update = getTimestamp()

            self._route_tablet = codec.decodeRouteTablet( self._route_tablet_text )
        ## updateRouteTable()

    def demandRoute(self, table_name, key) :
        self._logger.debug('demand route ...')

        try :
            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)
            s.connect( (self._meta_host, self._meta_port) )
        except :
            self._logger.excpt()
            self.onEndUpdateRouteTable()
            return

        request = Event()
        request.request_id = self.genRequestID()
        request.op = ROUTE_DEMAND
        request.body = (table_name, key)

        stream = iostream.IOStream(s, io_loop=self._io_loop)
        MetaConnection(stream, (self._meta_host, self._meta_port), request, self.demandRouteResponse)
        ## demandRoute()

    def demandRouteResponse(self, response) :
        if not response : return
        self._logger.debug('handle ROUTE_DEMAND response : %s', response)

        # update route table as soon as possible
        self._when_to_update_route_table = 0.5
        ## demandRouteResponse()

    ####-----------------------------------------------------------####

    def lookupAgent(self, packet) :
        ts_list = self.__lookupTabletServer(packet.table, packet.key)
        if not ts_list :
            if packet.op not in (Packet.OP_FIND,) :
                self.demandRoute(packet.table, packet.key)
            return None
 
        agents = []
        for ts in ts_list :
            if ts not in self._agents :
                q = queue.Queue()
                agent = Agent(self, ts, q, self._io_loop)
                self._agents[ ts ] = agent
            agents.append( self._agents[ ts ] )
       
        return agents
        ## lookupAgent()

    def __lookupTabletServer(self, table_name, key) :
        if table_name not in self._route_tablet :
            return None

        tablets = self._route_tablet[ table_name ]
        index = bisect_left(tablets, (key,key,None,None))
        if index > 0 : 
            index -= 1

        key_start,key_stop,ts_list = tablets[ index ]
        if (key_start < key) and (key <= key_stop) :
            return ts_list

        return None
        ## __lookupTabletServer()


    ####-----------------------------------------------------------####

    def genRequestID(self) :
        # TODO: lock
        serial = self._id_serial

        self._id_serial += 1
        if self._id_serial >= self._id_maximum :
            self._id_serial = 1

        id = '.'.join((str(self._id_base), str(serial)))
        return id
        ## genRequestID()


    ## class AgentPool

class   Agent(object) :
    def __init__(self, agent_pool, tablet_server_address, incoming_queue, io_loop) :
        self._q = incoming_queue
        self._ts_address = tablet_server_address
        self._agent_pool = agent_pool
        self._io_loop = io_loop

        self._ongoing = False
        self._keep_alive = True

        self._stream = None

        self._END_TAG = "\r\n\r\n"

        self._client_connection = None
        self._logger = logging.getLogger('gate.agent')
        ## __init__()

    def sendRequest(self, client_connection, request, timestamp, evt) :
        self._q.put( (client_connection, request, timestamp, evt) )

        self.__handleRequest()
        ## sendRequest()

    def __handleRequest(self) :
        if self._ongoing and self._stream :
            return

        if not self._ongoing :
            try :
                _connection,_request,_timestamp,_event = self._q.get(False)
            except queue.Empty :
                return

            self._ongoing = True

            self._tag_ongoing = (_connection, _request, _timestamp)
            self._wait_response_for = _event.request_id
            request_str = codec.encode( _event )
            self._request_ongoing = request_str

        ok = self.__connect()
        if not ok :
            return

        self._stream.write( ''.join((self._request_ongoing, self._END_TAG)) )
        self._stream.read_until(self._END_TAG, self.__onResponse)
        ## __handleRequest()    

    def __onResponse(self, response_str) :
        self._q.task_done()

        response = codec.decode( response_str )
        if response.response_to == self._wait_response_for :
            # put response to request pool
            _connection, _request, _timestamp = self._tag_ongoing
            self._agent_pool.feedResponse(_connection, _request, _timestamp, self, response)

            # handle any queued request
            self._ongoing = False
            self._tag_ongoing = None
            self._request_ongoing = None
            self.__handleRequest()
        else :
            # TODO : ignore other response and request
            self._stream.read_until(self._END_TAG, self.__onResponse)
        ## __onResponse()

    def __connect(self) :
        if self._stream :
            return True

        try :
            self._logger.debug('try to connect to %s ...', self._ts_address)
            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)
            s.connect( self._ts_address )
            self._logger.info('connect to %s done.', self._ts_address)
        except :
            s.close()
            # fail to connect
            self._logger.excpt()

            deadline = time.time() + 1.0
            self._io_loop.add_timeout(deadline, self.__tryToConnect)
            return False

        self._stream = iostream.IOStream( s, self._io_loop )
        self._stream.set_close_callback( self.__onDisconnect )
        return True
        ## __connect()

    def __disconnect(self, force=False) :
        if self._keep_alive and not force :
            return

        self._stream.close()
        self._stream = None
        ## __disconnect()

    def __onDisconnect(self) :
        self._logger.error('ERROR : connection to %s broke!', self._ts_address)
        self.__disconnect( force=True )

        deadline = time.time() + 1.0
        self._io_loop.add_timeout(deadline, self.__tryToConnect)
        ## __onDisconnect()

    def __tryToConnect(self) :
        if self._stream :
            return

        self.__handleRequest()
        ## __tryToConnect()

    ## class Agent

class   MetaConnection(object) :
    def __init__(self, stream, address, request, response_callback, end_callback=None) :
        self._stream = stream
        self._address = address
        self._response_callback = response_callback
        self._end_callback = end_callback

        self._keep_alive = False

        self._END_TAG = "\r\n\r\n"
        self._request_finished = False

        self._logger = logging.getLogger('metaconnect')
        self._logger.debug('create MetaConnection for %s', self._address)

        self.__wait_response_for = request.request_id
        request_str = codec.encode(request)
        self._stream.write( ''.join((request_str, self._END_TAG)) )
        self._stream.read_until(self._END_TAG, self.onResponse)
        ## __init__()

    def __del__(self) :
        #if self._stream.closed() :
        #    self._logger.debug('stream is closed when destroying')
        #else :
        #    self._logger.debug('stream is opened when destroying')
        self._logger.debug('destroy MetaConnection for %s', self._address)
        pass
        ## __del__()

    def onResponse(self, response_str) :
        response = codec.decode( response_str )
        if response.response_to == self.__wait_response_for :
            self._response_callback( response )

            self._stream.close()
            if self._end_callback : self._end_callback()
        else :
            # TODO : ignore other response and request
            self._stream.read_until(self._END_TAG, self.onResponse)
        ## onResponse()


    ## class MetaConnection

#######################################################################


########################################################################

def run(tcp_port, meta_url) :
    logger = logging.getLogger('gate')

    logger.info('gate startup (meta-server: %s) ...', meta_url)
    logger.info('  listening tcp port :%d ...', tcp_port)
    server = ServerThread(meta_url, tcp_port)
    server.start()

    try :
        while True :
            time.sleep(1.0)
    finally :
        logger.info('shutdown ...')
        server.stop()
        logger.info('joining server-thread ...')
        if server.isAlive() :
            server.join()

    logger.info('gate done')
    ## run()

if  __name__ == '__main__' :

    from optparse import OptionParser
    parser = OptionParser()
    parser.add_option("-m", "--meta", dest="meta_server", type="string", 
                        help='url of meta-server, e.g. tcp://192.168.1.2:4344')
    parser.add_option("-p", "--port", dest="port", type="int", default=14390,
                        help='listening tcp-port of gate')

    parser.add_option("-l", "--logconf", dest="logconfig", type="string",
                        help='path to logging configure file')

    parser.add_option("-d", "--debug", dest="debug", action="store_true", default=False,
                        help='show debug information')
    parser.add_option("-v", "--verbose", dest="verbose", action="store_true", default=False,
                        help='show detail information')

    (options, args) = parser.parse_args()
    if not options.meta_server :
        parser.print_help( sys.stdout )
        sys.exit(0)

    if options.logconfig :
        logging.loadConfigure( options.logconfig )

    if options.debug :
        logger = logging.getLogger()
        logger.setLevel(logging.LEVEL_DEBUG)

    run(options.port, options.meta_server)

